from flask import Flask, request, send_from_directory,jsonify

from threading import Thread
import time
import os

app = Flask(__name__)

@app.route('/start_training', methods=['POST'])
def start_training():
    task_id = request.json.get('taskid')
    if not task_id:
        return jsonify({"error": "缺失 taskid"}), 400

    # 创建训练任务的目录
    training_dir = os.path.join('training_tasks', task_id)
    if not os.path.exists(training_dir):
        os.makedirs(training_dir)

    # 启动一个新的线程来执行训练函数
    training_thread = Thread(target=perform_training, args=(task_id, training_dir))
    training_thread.start()

    return jsonify({"message": f"ai 训练服务已启动 {task_id}  "}), 200

def perform_training(task_id, training_dir):
    # 模拟训练过程
    print(f"Simulating training for task {task_id}...")

    # 创建一个空的txt文件作为训练结果
    with open(os.path.join(training_dir, 'training_result.txt'), 'w') as f:
        f.write(f"这是task id {task_id}的日志.\n")

    # 模拟耗时操作
    time.sleep(2)

    print(f"完成Training for task {task_id} completed.")
# 新增：下载预测结果文件
@app.route('/get_prediction/<output_dir>', methods=['GET'])
def get_prediction(output_dir):
    prediction_file_path = os.path.join('download', output_dir, 'ai小美满.wav')
    if os.path.exists(prediction_file_path):
        return send_from_directory(os.path.dirname(prediction_file_path), os.path.basename(prediction_file_path), as_attachment=True)
    else:
        return jsonify({"error": "Prediction result not found"}), 404
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6006)